#include "IOUring.h"
#include "IOUringEventLoop.h"
#include "Logger.h"
#include <string.h>

IOUring::IOUring(int entries, int groupId, IOUringEventLoop *loop)
    : loop_(loop), groupId_(groupId), entries_(entries) {
    struct io_uring_params params;
    memset(&params, 0, sizeof(params));
    if (io_uring_queue_init_params(entries_, &ring_, &params) < 0) {
        LOG_FATAL("io_uring init failed...");
    }
    if (!(params.features & IORING_FEAT_FAST_POLL)) {
        LOG_FATAL(
            "IORING_FEAT_FAST_POLL not available in the kernel, quiting...");
    }
    struct io_uring_probe *probe;
    probe = io_uring_get_probe_ring(&ring_);
    if (!probe ||
        !io_uring_opcode_supported(probe, IORING_OP_PROVIDE_BUFFERS)) {
        LOG_FATAL("Buffer select not supported, skipping...");
    }
    free(probe);
}

struct io_uring *IOUring::ring() { return &ring_; }

void IOUring::addProvideBuf(uint16_t bid) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
    io_uring_prep_provide_buffers(sqe, bufs_[bid], kMaxMessageLen, 1, groupId_,
                                  bid);
    conn_info conni = {.fd_ = 0, .type_ = PROV_BUF, .bid_ = (uint16_t)-1};
    memcpy(&sqe->user_data, &conni, sizeof(conni));
}
void IOUring::addProvideBufAll() {
    struct io_uring_sqe *sqe;
    struct io_uring_cqe *cqe;
    sqe = io_uring_get_sqe(&ring_);
    io_uring_prep_provide_buffers(sqe, bufs_, kMaxMessageLen, kBuffersCount,
                                  groupId_, 0);
    io_uring_submit(&ring_);
    io_uring_wait_cqe(&ring_, &cqe);
    if (cqe->res < 0) {
        printf("cqe->res = %d\n", cqe->res);
        exit(1);
    }
    io_uring_cqe_seen(&ring_, cqe);
}
void IOUring::addAccept(int fd, InetAddress &peerAddr, unsigned flags) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
    socklen_t clientLen = sizeof(*peerAddr.getSockAddr());
    io_uring_prep_accept(sqe, fd, (struct sockaddr *)peerAddr.getSockAddr(),
                         &clientLen, 0);
    io_uring_sqe_set_flags(sqe, flags);

    conn_info conni = {
        .fd_ = (uint32_t)fd, .type_ = ACCEPT, .bid_ = (uint16_t)-1};
    memcpy(&sqe->user_data, &conni, sizeof(conni));
}

void IOUring::addRead(int fd, size_t size, uint32_t flags) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
    io_uring_prep_recv(sqe, fd, NULL, size, 0);
    io_uring_sqe_set_flags(sqe, flags);
    sqe->buf_group = groupId_;

    conn_info conni = {
        .fd_ = (uint32_t)fd, .type_ = READ, .bid_ = (uint16_t)-1};
    memcpy(&sqe->user_data, &conni, sizeof(conni));
}

void IOUring::addWrite(int fd, uint16_t bid, size_t size, uint32_t flags) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
    io_uring_prep_send(sqe, fd, &bufs_[bid], size, 0);
    io_uring_sqe_set_flags(sqe, flags);

    conn_info conni = {
        .fd_ = (uint32_t)fd,
        .type_ = WRITE,
        .bid_ = bid,
    };
    memcpy(&sqe->user_data, &conni, sizeof(conni));
}

void IOUring::handleCqes() {
    uint32_t head = 0;
    uint32_t count = 0;
    struct io_uring_cqe *cqe;
    io_uring_for_each_cqe(&ring_, head, cqe) {
        ++count;
        struct conn_info conni;
        memcpy(&conni, &cqe->user_data, sizeof(conni));

        int type = conni.type_;
        if (cqe->res == -ENOBUFS) { // 没有空余的 buf
            LOG_FATAL("bufs in automatic buffer selection empty, this should "
                      "not happen...");
        } else if (type == PROV_BUF) { // 添加 buf
            if (cqe->res < 0) {
                LOG_FATAL("cqe->res = %d", cqe->res);
            }
        } else if (type == ACCEPT) {
            loop_->hasPrepAccept_ = false;
            LOG_DEBUG("accept %d %s", loop_->acceptFd_,
                      loop_->peerAddr_.toIpPort().c_str());
            if (newConnectionCallback_) {
                newConnectionCallback_(cqe);
            }
        } else if (type == READ) {
            if (readCompleteCallback_) {
                readCompleteCallback_(cqe);
            }
        } else if (type == WRITE) {
            if (writeCompleteCallback_) {
                writeCompleteCallback_(cqe);
            }
        }
    }

    io_uring_cq_advance(&ring_, count);
}
